#!/usr/bin/python
# -*- coding: UTF-8 -*-

import os
import xml.dom.minidom


def findfiles(path, t):
    '''
    遍历查找目录下的所有jpg文件名
    :param path:
    :param t:
    :return:
    '''
    files = os.listdir(path)
    for f in files:
        npath = path + '/' + f
        if(os.path.isfile(npath)):
            if(os.path.splitext(npath)[1] ==".jpg"):
                t.append(os.path.splitext(f)[0])
    return

def xmltotxt(inputfolder_xml, filename, inputfolder_img, classes):
    '''
    对指定目录下的指定名字的xml文件尽心格式转换，并保存到指定目录下
    :param inputfolder_xml:
    :param filename:
    :param inputfolder_img:
    :return:vi
    '''

    #打开标注数据xml文件
    xmlpath = inputfolder_xml + '/' + filename + '.xml'
    DOMTree = xml.dom.minidom.parse(xmlpath)
    collection = DOMTree.documentElement
    W = float(collection.getElementsByTagName("width")[0].childNodes[0].data)
    H = float(collection.getElementsByTagName("height")[0].childNodes[0].data)

    skuname = collection.getElementsByTagName("name")
    skudifficult = collection.getElementsByTagName("difficult")
    xmlxmin = collection.getElementsByTagName("xmin")
    xmlxmax = collection.getElementsByTagName("xmax")
    xmlymin = collection.getElementsByTagName("ymin")
    xmlymax = collection.getElementsByTagName("ymax")
    rectNum = len(xmlxmin)
    #print("")

    #将格式转换后保存至txt文件
    txtpath = outputfolder_txt + "/" + filename + ".txt"
    f1 = open(txtpath, 'w')
    for i in range(rectNum):
        name = skuname[i].childNodes[0].data
        difficult = int(skudifficult[i].childNodes[0].data)
        if 1 == difficult:
            continue
        if name not in classes:
            print("sku name not in list: " + name + "int img " + filename)
            continue

        xmin = float(xmlxmin[i].childNodes[0].data) / W
        xmax = float(xmlxmax[i].childNodes[0].data) / W
        ymin = float(xmlymin[i].childNodes[0].data) / H
        ymax = float(xmlymax[i].childNodes[0].data) / H

        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2
        w = xmax - xmin
        h = ymax - ymin
        index = classes.index(name)

        f1.write(str(index) + " " + str(cx) + " " + str(cy) + " " + str(w) + " " + str(h))
        f1.write('\n')
    f1.close()


if __name__=="__main__":
    '''
    主函数
    输入：jpg图像目录，xml目录
    输出：将xml数据转换为txt格式并添加到图像目录中（txt文件数量与jpg相同，不一定与xml文件数量相同）
          同时在两目录外生成一个txt文件，记录所有jpg文件的路径    
    '''
    print("version:2018.11.14  author: cwh")
    wd = os.getcwd()

    errstatus = 0 #0无错误，-1有错误

    #sku列表
    classes = ['car', 'van', 'truck', 'bus']

    #输入输出路径
    inputfolder_img = "training_image_2000"
    inputfolder_xml = "training_image_2000_xml"#"/home/yunxing/Documents/dry_16sku_2018-11-22/traindataset"
    statisticTxtPath = "training_image_2000.txt"
    outputfolder_txt = "training_image_2000_txt"
    if not os.path.exists(outputfolder_txt):
        os.mkdir(outputfolder_txt)

    #检索所有jpg文件名
    jpgfilenamelist = []
    findfiles(inputfolder_img, jpgfilenamelist)
    print("found %d jpg files." % len(jpgfilenamelist))
    print("")

    #生成jpg文件路径列表文件，即yolo训练数据中的train.txt
    f1 = open(statisticTxtPath, 'w')
    for filename in jpgfilenamelist:
        jpgpath = inputfolder_img + "/" + filename + ".jpg"
        f1.write(jpgpath)
        f1.write('\n')
    f1.close()

    #将xml文件转为txt文件
    for filename in jpgfilenamelist:
        xmlpath = inputfolder_xml + "/" + filename + ".xml"
        #print("processing " + xmlpath)
        if (False == os.path.exists(xmlpath)):
            print("ERROR! %s's xml file does not exists!" % filename)
            errstatus = -1
            break
        else:
            #print("processing %s ..." % filename)
            xmltotxt(inputfolder_xml, filename, inputfolder_img, classes)

    #jpgfilenamelist.clear()

